import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from utils.Logger import get_logger

log = get_logger()

# 定义颜色列表
rain_list = [
    "#a5f38e", "#3ccd07", "#3ca500", "#02724e", "#3b7ddb", "#2b5cc2",
    "#1c3ba9", "#112c90", "#fb00fb", "#c811a9", "#86158a", "#eb4900",
]

# 创建颜色映射
cmap = mpl.colors.ListedColormap(rain_list)
cmap.set_bad(alpha=0.0)  # 设置缺失值颜色为透明
cmap.set_over("#820040")  # 设置超出范围的值颜色
cmap.set_under(alpha=0.0)  # 设置低于范围的值颜色为透明

# 设置颜色映射的边界
bounds = [0.1, 1, 2, 5, 10, 25, 50, 75, 100, 150, 200, 250, 300]
norm = mpl.colors.BoundaryNorm(bounds, cmap.N)


def draw_cnn_png(inputFrame, save_path):
    '''

    :param inputFrame:numpy (400,400) 输入没有经过归一化 0-300
    :param save_path:
    :return:
    '''
    log.info("保存降雨外推结果" + str(save_path))
    fig = plt.figure(figsize=(1, 1))
    axes = fig.add_subplot(111)
    mask = (inputFrame < 0.1)
    data = np.ma.masked_where(mask, inputFrame)
    plt.contourf(data, levels=np.array(bounds + [500]), cmap=cmap, norm=norm)
    axes.set_axis_off()
    axes.invert_yaxis()
    plt.subplots_adjust(left=0, bottom=0, right=1, top=1, hspace=0.1, wspace=0.1)
    plt.savefig(save_path, dpi=400, format="png", transparent=True)
    plt.close('all')
    plt.clf()


